
import SimpleITK as sitk
import numpy as np 

from models.nnunet3d import get_nnunet3d
from monai import transforms

import torch 

import time 

import numpy as np 
import os 
# 
# os.environ["CUDA_VISIBLE_DEVICES"] = "2"

import torch
# 限制0号设备的显存的使用量为0.5，就是半张卡那么多，比如12G卡，设置0.5就是6G。
torch.cuda.set_per_process_memory_fraction(0.5, 2)
torch.cuda.empty_cache()
# 计算一下总内存有多少。
total_memory = torch.cuda.get_device_properties(2).total_memory
# 使用0.499的显存:
tmp_tensor = torch.empty(int(total_memory * 0.499), dtype=torch.int8, device='cuda:2')

# 清空该显存：
del tmp_tensor
torch.cuda.empty_cache()

# 下面这句话会触发显存OOM错误，因为刚好触碰到了上限:
torch.empty(total_memory // 3, dtype=torch.int8, device='cuda:2')



# s = time.time()
# data = np.load("./data/fullres/train/segrap_0000.npy", "r")
# e = time.time()
# print(f"spend time: {e - s}")
# print(data.shape)

# with torch.no_grad():
#     t1 =  torch.zeros((45, 128, 450, 960))
#     print(f"构建成功")
#     time.sleep(10)

#     t1 =  torch.zeros((45, 128, 450, 960), dtype=torch.int8)
#     # t1 =  torch.zeros((45, 128, 450, 960))
#     print(f"构建成功")
#     t1 = t1.to("cuda:3")

#     time.sleep(10)


# import torch 
# from models.modelgenesis.unet3d import UNet3DModelGen

# model = UNet3DModelGen(2, 45)

# t1 = torch.rand(1, 2, 64, 256, 256)

# out = model(t1)

# for o in out:
#     print(o.shape)

# t1 = torch.rand(1, 2, 64, 256, 256)

# t2 = torch.rand(1, 64, 256, 256)

# t1[:, 0] = t2

# model = get_nnunet3d(2, 54)
# out = model(t1)
# for o in out:
#     print(o.shape)
    
# # label_1 = "/home/xingzhaohu/jiuding_code/SegRap2023/SegRap2023_Training_Set_120cases/segrap_0000/Brain.nii.gz"
# # label_1 = "/home/xingzhaohu/jiuding_code/SegRap2023/SegRap2023_Training_Set_120cases/segrap_0000/Brain.nii.gz"
# label_1 = "/home/xingzhaohu/jiuding_code/SegRap2023/data/raw_data/SegRap2023_Training_Set_120cases/segrap_0000/seg.nii.gz"
# label_1 = "/home/xingzhaohu/jiuding_code/SegRap2023/SegRap2023_Training_Set_120cases_OneHot_Labels/Task002/segrap_0002.nii.gz"
# label_1 = sitk.ReadImage(label_1)
# label_1 = sitk.GetArrayFromImage(label_1)

# print(np.unique(label_1))
# # image_1 = "/home/xingzhaohu/jiuding_code/SegRap2023/SegRap2023_Training_Set_120cases/segrap_0114/image.nii.gz"
# # image_1 = sitk.ReadImage(image_1)
# # image_1 = sitk.GetArrayFromImage(image_1)
# # image_2 = "/home/xingzhaohu/jiuding_code/SegRap2023/SegRap2023_Training_Set_120cases/segrap_0114/image_contrast.nii.gz"
# # image_2 = sitk.ReadImage(image_2)
# # image_2 = sitk.GetArrayFromImage(image_2)
# print(np.sum(label_1))


# data = "/home/xingzhaohu/jiuding_code/SegRap2023/data/fullres/train/segrap_0000.npz"
# data = np.load(data)
# image = data["data"]
# label_2 = data["seg"]

# label_2[label_2 == -1] = 0
# print(np.sum(label_2))

# import matplotlib.pyplot as plt 

# # print(image_1.shape)
# # print(np.unique(image_1))
# # print(np.unique(image_2))
# for i in range(80, 100):
#     plt.subplot(1, 3, 1)
#     plt.imshow(image[0, i], cmap="gray")
#     # plt.subplot(1, 3, 2)
#     # plt.imshow(label_1[i], cmap="gray")
#     plt.subplot(1, 3, 2)
#     plt.imshow(label_2[0, i], cmap="gray")
#     plt.show()